# -*- coding: utf-8 -*-
# Author   : ZhangQing
# Time     : 2025-07-08 6:47
# File     : integration_tests.py
# Project  : risk-contagion-analysis
# Desc     :

import unittest
import pandas as pd
import numpy as np
from src.rust_bridge import SupplyChainGraph, RiskContagionSimulator, NetworkAnalyzer
from src.data_processor import DataProcessor
from src.visualization import NetworkVisualizer


class TestFullPipeline(unittest.TestCase):
    def setUp(self):
        # 创建测试数据
        companies_data = {
            'id': [f'C{i:03d}' for i in range(1, 21)],
            'name': [f'Company {i}' for i in range(1, 21)],
            'risk_score': np.random.uniform(0, 1, 20),
            'industry': np.random.choice(['Tech', 'Manufacturing', 'Retail', 'Finance'], 20),
            'size': np.random.randint(50, 500, 20)
        }

        # 创建关系数据
        relationships = []
        for i in range(1, 18):
            for j in range(i + 1, min(i + 4, 21)):
                if np.random.random() > 0.3:  # 70% 概率创建连接
                    relationships.append({
                        'source_id': f'C{i:03d}',
                        'target_id': f'C{j:03d}',
                        'strength': np.random.uniform(0.1, 0.9),
                        'relationship_type': np.random.choice(['supplier', 'customer', 'partner'])
                    })

        self.companies_df = pd.DataFrame(companies_data)
        self.relationships_df = pd.DataFrame(relationships)

    def test_end_to_end_pipeline(self):
        """测试完整的分析流程"""

        # 1. 数据处理
        processor = DataProcessor()
        processed_companies = processor.process_companies(self.companies_df)
        processed_relationships = processor.process_relationships(self.relationships_df)

        # 2. 构建网络
        graph = SupplyChainGraph()
        graph.build_from_dataframe(processed_companies, processed_relationships)

        # 3. 运行风险模拟
        simulator = RiskContagionSimulator()
        results = simulator.simulate(
            graph,
            contagion_factor=0.2,
            damping_factor=0.8
        )

        # 4. 进行网络分析
        analyzer = NetworkAnalyzer()
        centrality = analyzer.calculate_centrality(graph)

        # 5. 创建可视化（仅测试接口）
        visualizer = NetworkVisualizer()
        # 只测试方法能否执行，不测试图像输出
        try:
            visualizer.plot_risk_contagion(graph, results)
            visualizer.plot_centrality(graph, centrality)
            visualization_works = True
        except Exception as e:
            visualization_works = False
            print(f"可视化测试失败: {e}")

        # 验证整个流程能否顺利完成
        self.assertTrue(visualization_works)
        self.assertIsInstance(results, dict)
        self.assertIsInstance(centrality, dict)

    # 更多集成测试...


if __name__ == "__main__":
    unittest.main()
